
#include <string.h>
#include "../iface/CryptInterface.h"
#include "../iface/debug.h"
#include "../xml/DOMDocument.h"
#include "../utils/BaseConvert.h"
#include "../utils/StringUtils.h"
#include "../crypt/RSAKeyPair.h"
#include "SSHChain.h"

// By key or by password

SSHChain::SSHChain(ClusterSSHDB *iclusters, apr_pool_t *imp)
{
   apr_status_t status;
   status = apr_pool_create(&mp, imp);
   CheckAPRError(status);
   processed_for_unknown_fingerprints = 0;
   root = NULL;
   sso_keyname = NULL;
   sso_keypass = NULL;
   sso_userid = NULL;
   our_pass = NULL;
   clusterdb = iclusters;
}

// The idea behind this is that if the key pool does not have access to a needed resource, it can be added as this function will return the node required
char *SSHChain::BuildChainToCluster(const char *clusterxmlfile, const char *clusterid, const char *isso_userid, apr_pool_t *imp)
{
   sso_keyname = NULL;
   sso_keypass = NULL;
   if (sso_userid) free((void *)sso_userid);
   sso_userid = NULL;
   char *tmp_userid = (char *)malloc(sizeof(char)*(strlen(isso_userid)+1));
   strcpy(tmp_userid, isso_userid);
   sso_userid = tmp_userid;
   our_pass = NULL;
   processed_for_unknown_fingerprints = 0;
   DOMDocument *domd = new DOMDocument(clusterxmlfile);
   xmlNodePtr tnode = domd->findNode(domd->getRootNode(), "cluster", "idname", clusterid);
   root = NULL;
   xmlNodePtr cnode = tnode->children;
   while (strcmp((const char *)(cnode->name), "gatehop") == 0)
   {
      char *t_host = domd->getAttribValueString(cnode, "host");
      char *loginname = clusterdb->GetLogInNameForHostForUser(sso_userid, t_host);
      int hostokay = (loginname != NULL);
      free(loginname);
      if (hostokay)
      {
         char *t_port = domd->getAttribValueString(cnode, "port");
         char *t_host = domd->getAttribValueString(cnode, "host");
         char *t_user = clusterdb->GetLogInNameForHostForUser(sso_userid, t_host);
         char *t_fingerprint = domd->getAttribValueString(cnode, "fingerprint");
         int t_usekey = clusterdb->BooleanCanLogInWithKey(sso_userid, t_host);
         AddToSSHChain(t_user, (t_port) ? atoi(t_port) : 0, t_host, t_fingerprint, t_usekey);
         if (t_user) free((void *)t_user);
         if (t_port) free((void *)t_port);
         if (t_fingerprint) free((void *)t_fingerprint);
         if (!t_usekey) { delete domd; return t_host; }
         if (t_host) free((void *)t_host);
         cnode = cnode->children;
      } else
      {
         cnode = cnode->next;
      }
   }
   // last one is always 'finalhop' so
   cnode = tnode->children;
   while (strcmp((const char *)(cnode->name), "gatehop") == 0) cnode = cnode->next;
   if (strcmp((const char *)(cnode->name), "finalhop") == 0)
   {
      char *t_port = domd->getAttribValueString(cnode, "port");
      char *t_host = domd->getAttribValueString(cnode, "host");
      char *t_fingerprint = domd->getAttribValueString(cnode, "fingerprint");
      char *t_user = clusterdb->GetLogInNameForHostForUser(sso_userid, t_host);
      int t_usekey = clusterdb->BooleanCanLogInWithKey(sso_userid, t_host);
      AddToSSHChain(t_user, (t_port) ? atoi(t_port) : 0, t_host, t_fingerprint, t_usekey);
      if (t_user) free((void *)t_user);
      if (t_port) free((void *)t_port);
      if (t_host) free((void *)t_host);
      if (t_fingerprint) free((void *)t_fingerprint);
      if (!t_usekey) { delete domd; return t_host; }
      if (t_host) free((void *)t_host);
   } else
   {
      errprint("last xml tag is not 'finalhop'");
   }
   delete domd;
   return NULL;
}

void SSHChain::DustForFingerprints(PendingDB *pdb)
{
   SSHLink *tthis = root;
   while (tthis)
   {
      if ((tthis->obtained_host_fingerprint) && (strcmp(tthis->expected_host_fingerprint, tthis->obtained_host_fingerprint) != 0))
      {
         pdb->AddNewHostKeyUpdateRequest(tthis->remote_hostname,
            ((tthis->expected_host_fingerprint == NULL) || (strlen(tthis->expected_host_fingerprint) == 0)) ? "(none)" : tthis->expected_host_fingerprint,
            tthis->obtained_host_fingerprint);
      }
      tthis = tthis->next;
   }
   processed_for_unknown_fingerprints = 1;
}

void SSHChain::SetAuthenticatorKeys(KeyDB *sshkeys, const char *keyname, const char *keypass)
{
   keydb = sshkeys;
   char *tmp_keyname = (char *)malloc(sizeof(char)*(strlen(keyname)+1));
   strcpy(tmp_keyname, keyname);
   sso_keyname = tmp_keyname;
   char *tmp_keypass = (char *)malloc(sizeof(char)*(strlen(keypass)+1));
   strcpy(tmp_keypass, keypass);
   sso_keypass = tmp_keypass;
}

void SSHChain::SetAuthenticatorPassword(const char *authpass)
{
   char *tmp_pass = (char *)malloc(sizeof(char)*(strlen(authpass)+1));
   strcpy(tmp_pass, authpass);
   our_pass = (const char *)tmp_pass;
}

SSHChain::~SSHChain()
{
   DestroySSHChain(root);
   // Add in some memset 0 blanking for keys and passwords
   if (our_pass)
   {
      int ourpass_len = strlen(our_pass);
      memset((void *)our_pass, 0, sizeof(char)*ourpass_len);
      free((void *)our_pass);
   }
   if (sso_keyname)
   {
      int ssokeyname_len = strlen(sso_keyname);
      memset((void *)sso_keyname, 0, sizeof(char)*ssokeyname_len);
      free((void *)sso_keyname);
   }
   if (sso_keypass)
   {
      int ssokeypass_len = strlen(sso_keypass);
      memset((void *)sso_keypass, 0, sizeof(char)*ssokeypass_len);
      free((void *)sso_keypass);
   }
   if (sso_userid) free((void *)sso_userid);
   keydb = NULL;
   root = NULL;
   if (!processed_for_unknown_fingerprints)
      errprint("this SSH chain has not been processed for any unknown hosts it could encounter (forgot to dust for fingerprints?)");
   apr_pool_destroy(mp);
}

/*
   const char *remote_hostname;
   int remote_hostport;
   const char *real_username; // this is NULL if it is the final link
   int local_hostport; // set this to zero, if the remote hostport is to be used
   const char *key_fingerprint;
   struct ssh_link *next;
*/

SSHLink **SSHChain::GetLastChainLink()
{
   SSHLink *final_link = root;
   SSHLink **new_link = NULL;
   if (final_link != NULL)
   {
      while (final_link->next != NULL)
         final_link = final_link->next;
      new_link = &(final_link->next);
   } else new_link = &root;
   return new_link;
}

void SSHChain::AddToSSHChain(const char *isso_userid, const char *hostname)
{
   if (sso_userid) free((void *)sso_userid);
   sso_userid = NULL;
   char *tmp_userid = (char *)malloc(sizeof(char)*(strlen(isso_userid)+1));
   strcpy(tmp_userid, isso_userid);
   sso_userid = tmp_userid;
   char *unaliased_hostname = clusterdb->UnaliasHostName(hostname);
   std::vector<char *> *output = clusterdb->sqlScanAllTextf("select port,hostfingerprint from 'ssh_hosts' where hostname='%s';", unaliased_hostname);
   if ((output) && (output->size() > 0))
   {
      char *t_user = clusterdb->GetLogInNameForHostForUser(sso_userid, unaliased_hostname);
      int t_port = atoi(output->at(0));
      char *t_fingerprint = output->at(1);
      //printf("SUPPLIED FINGERPRINT = %s ofr %s\n", output->at(1), unaliased_hostname);
      int t_usekey = clusterdb->BooleanCanLogInWithKey(sso_userid, unaliased_hostname);
      AddToSSHChain(t_user, t_port, unaliased_hostname, t_fingerprint, t_usekey);
      free(t_user);
      StringUtils::destroyStringList(output);
   } else
   {
      if (output) delete output;
      errprint("could not find '%s' (%s) in 'ssh_hosts'", hostname, unaliased_hostname);
   }
   free(unaliased_hostname);
}

void SSHChain::AddToSSHChain(const char *user, int port, const char *host, const char *fingerprint, int usekey)
{
   SSHLink *tnlink = (SSHLink *)malloc(sizeof(SSHLink));
   tnlink->next = NULL;
   if (host)
   {
      char *tmp_hostname = clusterdb->UnaliasHostName(host);
      strcpy(tmp_hostname, host);
      tnlink->remote_hostname = tmp_hostname;
   } else tnlink->remote_hostname = NULL;
   tnlink->remote_hostport = port;
   tnlink->local_hostport = 0;
   if (user)
   {
      char *tmp_username = (char *)malloc(sizeof(char)*(strlen(user)+1));
      strcpy(tmp_username, user);
      tnlink->real_username = tmp_username;
   } else tnlink->real_username = NULL;
   if (fingerprint)
   {
      char *tmp_fingerprint = (char *)malloc(sizeof(char)*(strlen(fingerprint)+1));
      strcpy(tmp_fingerprint, fingerprint);
      tnlink->expected_host_fingerprint = tmp_fingerprint;
   } else tnlink->expected_host_fingerprint = NULL;
   tnlink->obtained_host_fingerprint = NULL;
   tnlink->usekey = usekey;
   tnlink->activated = 0;
   tnlink->scli = NULL;
   tnlink->fwds = NULL;
   SSHLink **new_link = GetLastChainLink();
   (*new_link) = tnlink;
}

void SSHChain::DestroySSHChain(SSHLink *root)
{
   if (root->next) DestroySSHChain(root->next);
   if (root->remote_hostname) free((void *)root->remote_hostname);
   if (root->real_username) free((void *)root->real_username);
   if (root->expected_host_fingerprint) free((void *)root->expected_host_fingerprint);
   if (root->obtained_host_fingerprint) free((void *)root->obtained_host_fingerprint);
   if (root->fwds) delete root->fwds;
   if (root->scli) delete root->scli;
   free(root);
}

char *SSHChain::GetFailedLinkHostName()
{
   SSHLink *tlink = root;
   while (tlink != NULL)
   {
      if (!(tlink->activated))
      {
         char *failedname = (char *)malloc(sizeof(char)*(strlen(tlink->remote_hostname)+1));
         strcpy(failedname, tlink->remote_hostname);
         return failedname;
      }
      tlink = tlink->next;
   }
   return NULL;
}

int SSHChain::WasFingerprintProblem()
{
   SSHLink *tlink = root;
   while (tlink != NULL)
   {
      if (!(tlink->activated))
      {
         if ((tlink->expected_host_fingerprint) && (tlink->obtained_host_fingerprint))
         {
            if (strcmp(tlink->expected_host_fingerprint, tlink->obtained_host_fingerprint) != 0)
            {
               return 1;
            }
         }
      }
      tlink = tlink->next;
   }
   return 0;
}

SSHClient *SSHChain::ActivateSSHChain(int *failcode)
{
   SSHLink *oldlink = NULL;
   SSHLink *tlink = root;
   int retval = 0;
   while ((tlink != NULL) && (!retval))
   {
      retval = ActivateSSHChainLink(tlink, (tlink->local_hostport == 0));
      oldlink = tlink;
      tlink = tlink->next;
   }
   if (failcode) *failcode = retval;
   return (oldlink) ? oldlink->scli : NULL;
}

int SSHChain::ActivateSSHChainLink(SSHLink *tlink, int useremote)
{
   const char *real_hostname = useremote ? tlink->remote_hostname : "localhost";
   if (!real_hostname)
   {
      errprint("hostname not specified");
      return 1;
   }
   const int real_hostport = useremote ? tlink->remote_hostport : tlink->local_hostport;
   if (!real_hostport)
   {
      errprint("host port not specified");
      return 1;
   }
   if (!(tlink->real_username))
   {
      errprint("username not specified");
      return 1;
   }
   SSHClient *ssh_cli = NULL;
   const char *hexfingerprint = NULL;
   int retries = 0;
   int maxretr = 10;
   do
   {
      ssh_cli = new SSHClient(real_hostname, real_hostport, mp);
      hexfingerprint = ssh_cli->DoHandshakeAndGetHostFingerprintInHex();
      if (hexfingerprint == NULL)
      {
         errprint("handshake try %d of %d", retries, maxretr);
         delete ssh_cli;
         ssh_cli = NULL;
         apr_sleep(rand() % 10000000);
      }
      retries++;
   } while ((hexfingerprint == NULL) && (retries < maxretr));
   tlink->obtained_host_fingerprint = hexfingerprint;
   if (tlink->obtained_host_fingerprint == NULL)
   {
      errprint("could not get host fingerprint");
      return 1;
   }
   if ((tlink->expected_host_fingerprint == NULL) || (strlen(tlink->expected_host_fingerprint) == 0))
   {
      errprint("no host fingerprint in database for host '%s', fingerprint was '%s'", tlink->remote_hostname, tlink->obtained_host_fingerprint);
      if (ssh_cli) delete ssh_cli;
      return 1;
   } else
   {
      if (strcmp(tlink->expected_host_fingerprint, tlink->obtained_host_fingerprint) != 0)
      {
         errprint("expected host fingerprint '%s', obtained fingerprint '%s'",
            ((tlink->expected_host_fingerprint == NULL) || (strlen(tlink->expected_host_fingerprint) == 0)) ? "(none)" : tlink->expected_host_fingerprint,
            tlink->obtained_host_fingerprint);
         if (ssh_cli) delete ssh_cli;
         return 1;
      }
   }
   if (tlink->usekey)
   {
      if (keydb && sso_keyname) // need both for SSH2!
      {
         char *lkeyname = clusterdb->CanLogInWithKey(sso_userid, tlink->remote_hostname);
         ssh_cli->SetAuthenticatorKeys(keydb, lkeyname, sso_keypass);
         free(lkeyname);
      } else
      {
         errprint("tried to use keys that don't exist!");
         delete ssh_cli;
         return 1;
      }
   } else
   {
      if (our_pass)
      {
         ssh_cli->SetAuthenticatorPassword(our_pass);
      } else if (keydb && sso_keyname)
      {
         char *lkeyname = clusterdb->CanLogInWithKey(sso_userid, tlink->remote_hostname);
         ssh_cli->SetAuthenticatorKeys(keydb, lkeyname, sso_keypass);
         free(lkeyname);
      } else
      {
         errprint("no authentication methods available for '%s'", tlink->remote_hostname);
         delete ssh_cli;
         return 1;
      }
   }
   int success;
   success = ssh_cli->SendUsernameAndDoAuthentication(tlink->real_username);
   char *sshkeyowner = StringUtils::varprintf(".ssh_%s", sso_userid);
   if (success != 0)
   {
      char *lkeyname = keydb->GetLatestPrivateKeyNameOfOwner(sshkeyowner);
      if (lkeyname)
      {
         ssh_cli->SetAuthenticatorKeys(keydb, lkeyname, sso_keypass);
         free(lkeyname);
         success = ssh_cli->ForceKeyAuthentication(tlink->real_username);
      }
      if (success != 0)
      {
         errprint("problem with authentication on host '%s'", tlink->remote_hostname);
         delete ssh_cli;
         return 1;
      }
   }
   const char *oldkeyname = clusterdb->CanLogInWithKey(sso_userid, tlink->remote_hostname);
   const char *newkeyname = keydb->GetLatestPrivateKeyNameOfOwner(sshkeyowner);
   // this key name is sso_keyname - no it's not anymore!
   if (newkeyname)
   {
      if ((oldkeyname == NULL) || (strcmp(oldkeyname, newkeyname) != 0))
      {
         RSAKeyPair *sshkey = (RSAKeyPair *)(keydb->GetPublicKeyByName(newkeyname));
         int keyadd_result = 1;
         for (int tries = 0; (keyadd_result == 1) && (tries < 10); tries++)
         {
            keyadd_result = ssh_cli->AddThisKey(newkeyname, sshkey);
            if (keyadd_result == 1)
            {
               errprint("key update try %d of %d failed", tries + 1, 10);
               apr_sleep(rand() % 1000000);
            }
         }
         if (keyadd_result == 0)
         {
            clusterdb->ChangeLogInKeyNameForUser(sso_userid, newkeyname, tlink->remote_hostname);
            if (oldkeyname) keydb->KeyNowUnusedByActor(oldkeyname);
            if (newkeyname) keydb->KeyPairNowUsedByActor(newkeyname);
            if ((oldkeyname) && (newkeyname)) printf("Updating from '%s' to '%s'", oldkeyname, newkeyname);
            // When the private key refcount reaches zero, it is an expired key
            std::vector<char *> *keynames = keydb->GetPrivateKeyNamesWithUsedCountAtZeroAndOwner(sshkeyowner);
            if (keynames)
            {
               // Each successful login goes through the list of expired but not dead keys and
               // makes sure that none of them are in the authorized_keys of the remote host
               for (int qa = 0; qa < keynames->size(); qa++)
               {
                  int keydelete_result = ssh_cli->DeleteThisKey(keynames->at(qa));
                  // if keydelete result is -1, error
                  // if keydelete result is 0, key deleted
                  // if keydelete result is 1, key not found to delete (so do nothing)
                  if (keydelete_result == 0)
                  {
                     errprint("deleting key '%s' from remote host '%s'", keynames->at(qa), tlink->remote_hostname);
                     keydb->KeyNowDeletedByActor(keynames->at(qa));
                  }
               }
               StringUtils::destroyStringList(keynames);
            }
         } else
         {
            errprint("key addition failed - probably won't be able to use node '%s'", tlink->remote_hostname);
         }
         delete sshkey;
      }
      if (oldkeyname) free((void *)oldkeyname);
      if (newkeyname) free((void *)newkeyname);
   }
   free(sshkeyowner);
   if (tlink->next)
   {
      int nosuccess = -1;
      int local_port_tunnel_from = -1;
      int triesleft = 3;
      ForwardServer *ifwds = NULL;
      while (nosuccess && triesleft)
      {
         local_port_tunnel_from = 22000 + (rand() % 1000);
         ifwds = new ForwardServer(local_port_tunnel_from, "127.0.0.1", mp);
         ifwds->StartSessionThreads();
         ifwds->StartListeningThread();
         // 1ms should be enough to ensure that the port is bound, if it is going to be
         apr_sleep(1000000);
         if (ifwds->IsListening())
         {
            nosuccess = ssh_cli->OpenTunnelChannel(tlink->next->remote_hostname, tlink->next->remote_hostport, local_port_tunnel_from);
            if (nosuccess)
               delete ifwds;
            else
               ifwds->SetForwardingChannel(ssh_cli->GetSSHClientChannel());
         } else
         {
            errprint("failed to bind to port %d", local_port_tunnel_from);
            nosuccess = 1;
            delete ifwds;
         }
         triesleft--;
      }
      if (!triesleft)
      {
         errprint("ran out of tries constructing port forwarding from '%s' to '%s'!", tlink->remote_hostname, tlink->next->remote_hostname);
         return 1;
      }
      tlink->fwds = ifwds;
      tlink->next->local_hostport = local_port_tunnel_from;
   }
   tlink->activated = 1;
   tlink->scli = ssh_cli;
   // Don't need to look for unknown fingerprints, all fingerprints were okay!
   processed_for_unknown_fingerprints = 1;
   return 0;
}
